jax==0.4.11
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jaxlib==0.4.7+cuda11.cudnn82
tensorflow>=2.6.0
tensorboard>=2.10.9
flax>=0.7.0
numpy>=1.24.1
pandas>=1.1.5
absl-py>=1.4.0
clu>=0.0.9
optax>=0.1.7
ml-collections>=0.1.1